package io.kumao.gateway.filter;

import com.google.common.collect.Lists;
import io.kumao.gateway.model.RateLimitConfig;
import io.kumao.gateway.repository.RateLimitConfigRepository;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.core.Ordered;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.PathContainer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.pattern.PathPatternParser;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.time.Instant;
import java.util.*;
import java.util.stream.Stream;

import static java.util.stream.Collectors.toList;
import static org.springframework.http.HttpHeaders.AUTHORIZATION;
import static org.springframework.http.server.PathContainer.parsePath;

/**
 * Redis限流全局过滤器.
 *
 * @author 崔作利
 */
@Slf4j
@Component
public class RedisRateLimitGlobalFilter implements GlobalFilter, Ordered {

    private final static String REMAINING_HEADER = "X-RateLimit-Remaining";
    private final static String REPLENISH_RATE_HEADER = "X-RateLimit-Replenish-Rate";
    private final static String BURST_CAPACITY_HEADER = "X-RateLimit-Burst-Capacity";
    private final static String REQUESTED_TOKENS_HEADER = "X-RateLimit-Requested-Tokens";

    private final RedisScript<List<Long>> script;
    private final PathPatternParser pathPatternParser = new PathPatternParser();

    private List<RateLimitConfig> configs = Lists.newArrayList();

    private boolean includeHeaders = true;

    @Resource
    private ReactiveStringRedisTemplate reactiveStringRedisTemplate;

    @Resource
    private RateLimitConfigRepository rateLimitConfigRepository;

    public RedisRateLimitGlobalFilter() {
        this.script = redisRequestRateLimiterScript();
    }

    @PostConstruct
    public void init() {
        List<RateLimitConfig> configs = rateLimitConfigRepository.selectList();
        log.info("限流过滤器 - 初始化配置({})", configs.size());
        for (RateLimitConfig config : configs) {
            log.info("限流过滤器 - 加载配置 {}", config);
            config.setPathPattern(pathPatternParser.parse(config.getPath()));
        }
        log.info("限流过滤器 - 初始化完成({})", configs.size());
        this.configs = configs;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest req = exchange.getRequest();
        ServerHttpResponse res = exchange.getResponse();

        HttpHeaders headers = req.getHeaders();
        List<String> tokens = headers.get(AUTHORIZATION);
        String token = tokens.get(0);

        if (CollectionUtils.isNotEmpty(tokens)) {
            PathContainer path = parsePath(req.getURI().getRawPath());

            Optional<RateLimitConfig> op = configs.stream().filter(config -> config.getPathPattern().matches(path)).findFirst();
            RateLimitConfig config = op.orElse(null);
            if (config != null) {
                return isAllowed(config, token).flatMap(response -> {
                    for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                        res.getHeaders().add(header.getKey(), header.getValue());
                    }

                    if (response.isAllowed()) {
                        return chain.filter(exchange);
                    }

                    log.warn("限流过滤器 - {} [{}] - {} [{}] / {}", ip(req), token, req.getURI(), config.getPath(), response.toString());
                    res.setStatusCode(config.getStatusCode());
                    return res.setComplete();
                });
            }
        }
        return chain.filter(exchange);
    }

    /**
     * 是否限流
     *
     * @param config
     * @param key
     * @return
     */
    private Mono<RateLimiter.Response> isAllowed(RateLimitConfig config, String key) {
        try {
            List<String> keys = getKeys(key);
            // The arguments to the LUA script. time() returns unixtime in seconds.
            List<String> scriptArgs = Arrays.asList(config.getReplenishRate() + "",
                    config.getBurstCapacity() + "", Instant.now().getEpochSecond() + "",
                    config.getRequestedTokens() + "");
            // allowed, tokens_left = redis.eval(SCRIPT, keys, args)
            Flux<List<Long>> flux = this.reactiveStringRedisTemplate.execute(this.script, keys, scriptArgs);
            // .log("redisratelimiter", Level.FINER);
            return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
                    .reduce(new ArrayList<Long>(), (longs, l) -> {
                        longs.addAll(l);
                        return longs;
                    }).map(results -> {
                        boolean allowed = results.get(0) == 1L;
                        Long tokensLeft = results.get(1);

                        return new RateLimiter.Response(allowed, getHeaders(config, tokensLeft));
                    });
        } catch (Exception e) {
            /*
             * We don't want a hard dependency on Redis to allow traffic. Make sure to set
             * an alert so you know if this is happening too much. Stripe's observed
             * failure rate is 0.01%.
             */
            log.error("Error determining if user allowed from redis", e);
        }
        return Mono.just(new RateLimiter.Response(true, getHeaders(config, -1L)));
    }

    @Override
    public int getOrder() {
        return -1;
    }

    public boolean isIncludeHeaders() {
        return includeHeaders;
    }

    public void setIncludeHeaders(boolean includeHeaders) {
        this.includeHeaders = includeHeaders;
    }

    private Map<String, String> getHeaders(RateLimitConfig config, Long tokensLeft) {
        Map<String, String> headers = new HashMap<>();
        if (isIncludeHeaders()) {
            headers.put(REMAINING_HEADER, tokensLeft.toString());
            headers.put(REPLENISH_RATE_HEADER, String.valueOf(config.getReplenishRate()));
            headers.put(BURST_CAPACITY_HEADER, String.valueOf(config.getBurstCapacity()));
            headers.put(REQUESTED_TOKENS_HEADER, String.valueOf(config.getRequestedTokens()));
        }
        return headers;
    }

    /**
     * Redis Lua 限流 Keys
     *
     * @param id
     * @return
     */
    private List<String> getKeys(String id) {
        // use `{}` around keys to use Redis Key hash tags
        // this allows for using redis cluster

        // Make a unique key per user.
        String prefix = "request_rate_limiter.{" + id;

        // You need two Redis keys for Token Bucket.
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    /**
     * Redis Lua 限流脚本
     *
     * @return
     */
    private DefaultRedisScript redisRequestRateLimiterScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/scripts/request_rate_limiter.lua")));
        redisScript.setResultType(List.class);
        return redisScript;
    }

    private String ip(ServerHttpRequest req) {
        HttpHeaders headers = req.getHeaders();
        List<String> headerIPs = headers.get("X-Forwarded-For");
        String ip = "";
        if (CollectionUtils.isNotEmpty(headerIPs)) {
            String headerIP = headerIPs.get(0);
            List<String> ips = Stream.of(headerIP.split(",")).filter((e) -> !StringUtils.startsWith(e, "100")).collect(toList());
            if (CollectionUtils.isNotEmpty(ips)) {
                ip = ips.get(0);
            }
        }
        return ip;
    }
}
